
from __future__ import division
from __future__ import print_function
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from joblib import dump
import keyboard
import os

from pyod.models.iforest import IForest
from pyod.models.knn import KNN
from pyod.models.pca import PCA
from pyod.utils.data import evaluate_print
from pyod.utils.example import visualize, data_visualize

if __name__ == "__main__":

    plt.switch_backend('TkAgg')
    columns_name = ["DtaNgx_1s", "DtaNgy_1s", "DtaNgz_1s", "DtaNgs_1s", "DtaNgt_1s"]
    columns_index = [37, 38, 39, 40, 41]

    # columns_name = ["DtaNax_1s", "DtaNay_1s", "DtaNaz_1s", "DtaNas_1s", "DtaNat_1s"]
    # columns_index = [42, 43, 44, 45, 46]

    file_names = os.listdir("./data/02.IMU") 

    data = None
    for file_name in file_names:
        if file_name.endswith(".dat"):
            file_path = os.path.join("./data/02.IMU", file_name)
            if data is None:
                data = pd.read_csv(file_path, sep='\t', usecols=columns_index, encoding="unicode_escape")
            else:
                temp_data = pd.read_csv(file_path, sep='\t', usecols=columns_index, encoding="unicode_escape")
                data = pd.concat([data, temp_data], ignore_index=True)
    
    print(data.head(5))

    data_len = data.shape[0]

    train_len = int(data_len * 0.9)
    test_len = int(data_len * 0.1)


    x_test = data.values[:test_len]
    x_train = data.values[test_len:]
    y_train = np.zeros(data_len - test_len,)
    y_test = np.zeros(test_len,)

    # 制造故障数据 陀螺
    x_train[10000] = [20, 30, 15, 17, 18]
    x_train[20000] = [0, 0, 0, 0, 0]
    x_train[30000] = [10, 13, 2, 8, 12]
    x_train[40000] = [14, 15, 7, -8, 15]

    # 制造故障数据 加表
    # x_train[10000] = [20, 30, 15, 17, 18]
    # x_train[20000] = [2400, -50, 4, -1400, 1690] 
    # x_train[30000] = [2314, 0, -1, -1413, 1600]
    # x_train[40000] = [2000, 0, -12, -1212, 1800]
    # x_train[50000] = [0, 0, 0, 0, 0]

    # plt.plot(x_train[...,0])
    # plt.show()

    # train IForest detector
    # clf_name = 'IForest'5
    # clf = IForest()
    # clf_name = 'KNN'
    # clf = KNN(contamination=0.00001)
    clf_name = "PCA"
    clf = PCA(contamination=0.000004, n_components=5) #contamination=0.000005 

    # clf = OCSVM() #contamination=0.000005
    clf.fit(x_train)

    # get the prediction labels and outlier scores of the training data
    y_train_pred = clf.labels_  # binary labels (0: inliers, 1: outliers)
    y_train_scores = clf.decision_scores_  # raw outlier scores

    print(np.count_nonzero(y_train_pred))
    print(np.nonzero(y_train_pred))
    print(x_train[np.nonzero(y_train_pred)])
    
    # get the prediction on the test data
    y_test_pred = clf.predict(x_test)  # outlier labels (0 or 1)
    y_test_scores = clf.decision_function(x_test)  # outlier scores
    print(np.count_nonzero(y_test_pred))


    # save model to disk
    dump(clf, './model_file/IMU_Gyro_Model.joblib')
    # dump(clf, './model_file/IMU_Acce_Model.joblib')

    print("Model saved successfully.")

    keyboard.wait("esc")
